/*
* Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

use std::{collections::{btree_map::Entry, BTreeMap}, path::Path};

use anyhow::Context;
use tree_sitter::{Node, Parser};

#[derive(Debug, PartialEq, Eq)]
pub struct MsgDef {
    pub fields: Vec<Field>,
}

#[derive(Debug, Default)]
pub struct MsgDefMap {
    pub map: BTreeMap<String, BTreeMap<String, MsgDef>>,
}

#[derive(Debug, PartialEq, Eq)]
pub struct Field {
    pub field_type: FieldType,
    pub field_name: String,
}

#[derive(Debug, PartialEq, Eq)]
pub struct FieldType {
    pub base_type: BaseType,
    pub array_size: Option<u32>,
    pub is_array: bool,
}

#[derive(Debug, PartialEq, Eq)]
pub struct BaseType {
    pub string: String,
    pub sep: Option<usize>,
}

#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum MsgDefParseError {
    #[error(transparent)]
    TreeSitterError(#[from] tree_sitter::LanguageError),
    #[error("Invalid format")]
    InvalidFormat,
    #[error("Invalid tree")]
    InvalidTree,
    #[error("No document node")]
    NoDocumentNode,
    #[error("Invalid array size: {0}")]
    InvalidArraySize(#[from] std::num::ParseIntError),
}

impl IntoIterator for MsgDefMap {
    type Item = (String, BTreeMap<String, MsgDef>);
    type IntoIter = std::collections::btree_map::IntoIter<String, BTreeMap<String, MsgDef>>;

    fn into_iter(self) -> Self::IntoIter {
        self.map.into_iter()
    }
}

impl MsgDefMap {
    pub fn new() -> Self {
        Self {
            map: BTreeMap::new(),
        }
    }

    pub fn iter(&self) -> std::collections::btree_map::Iter<String, BTreeMap<String, MsgDef>> {
        self.map.iter()
    }

    pub fn len(&self) -> usize {
        self.map.len()
    }

    pub async fn from_msgs_dir<P: AsRef<Path>>(msgs_dir: P) -> anyhow::Result<Self> {
        let msgs_dir = msgs_dir.as_ref();
        let mut msg_def_map = Self::new();
        let mut read_dir = tokio::fs::read_dir(msgs_dir).await
            .context(format!("failed to read '{}'", msgs_dir.display()))?;
        while let Some(entry) = read_dir.next_entry().await? {
            let file_type = entry.file_type().await
                .context(format!("failed to get file type of '{}'", entry.path().display()))?;
            if !file_type.is_dir() {
                continue;
            }

            let dir = entry.path();
            let mut read_dir = tokio::fs::read_dir(&dir).await
                .context(format!("failed to read '{}'", dir.display()))?;
            while let Some(entry) = read_dir.next_entry().await? {
                let file_type = entry.file_type().await
                    .context(format!("failed to get file type of '{}'", entry.path().display()))?;
                if !file_type.is_file() {
                    continue;
                }

                let file = entry.path();
                let message_string = tokio::fs::read_to_string(&file).await
                    .context(format!("failed to read '{}'", file.display()))?;
                msg_def_map.try_add_ros_msg_def(
                    dir.file_name().unwrap().to_string_lossy().to_string(),
                    file.file_stem().unwrap().to_string_lossy().to_string(),
                    &message_string,
                ).context(format!("failed to parse message definition in '{}'", file.display()))?;
            }
        }
        msg_def_map.resolve_field_types();
        Ok(msg_def_map)
    }

    pub fn try_add_ros_msg_def(
        &mut self,
        pkg_name: String,
        msg_name: String,
        message_string: &str,
    ) -> Result<bool, MsgDefParseError> {
        match self.map.entry(pkg_name) {
            Entry::Vacant(vacant_entry) => {
                let msg_def = Self::parse_ros_msg_def(message_string)?;
                let mut submap = BTreeMap::new();
                submap.insert(msg_name, msg_def);
                vacant_entry.insert(submap);
            },
            Entry::Occupied(mut occupied_entry) => {
                let submap = occupied_entry.get_mut();
                if submap.contains_key(&msg_name) {
                    return Ok(false);
                }
                let msg_def = Self::parse_ros_msg_def(message_string)?;
                submap.insert(msg_name, msg_def);
            },
        }

        Ok(true)
    }
    
    fn parse_ros_msg_def(content: &str) -> Result<MsgDef, MsgDefParseError> {
        let mut parser = Parser::new();
        parser.set_language(&tree_sitter_rosmsg::LANGUAGE.into())?;
        let tree = parser.parse(content.as_bytes(), None)
            .ok_or(MsgDefParseError::InvalidTree)?;
        let document_node = tree.root_node();
        if document_node.kind() != "document" {
            return Err(MsgDefParseError::NoDocumentNode);
        }
    
        let mut msg_def = MsgDef {
            fields: Vec::new(),
        };
        let mut maybe_node = document_node.named_child(0);
        while let Some(node) = maybe_node {
            if node.kind() == "definition" {
                let field_type_node = node.named_child(0)
                    .ok_or(MsgDefParseError::InvalidFormat)?;
                let field_name_node = node.named_child(1)
                    .ok_or(MsgDefParseError::InvalidFormat)?;
                debug_assert_eq!(field_type_node.kind(), "fieldtype");
    
                let type_kind_node = field_type_node.child(0)
                    .ok_or(MsgDefParseError::InvalidFormat)?;
                let field_name = field_name_node.utf8_text(content.as_bytes())
                    .map_err(|_| MsgDefParseError::InvalidFormat)?;
                match type_kind_node.kind() {
                    "basetype" => {
                        let base_type = Self::parse_base_type(content, type_kind_node)?;
                        let field_type = FieldType {
                            base_type,
                            array_size: None,
                            is_array: false,
                        };
                        msg_def.fields.push(Field {
                            field_type,
                            field_name: field_name.to_string(),
                        });
                    },
                    "arraytype" => {
                        let base_type_node = type_kind_node.named_child(0)
                            .ok_or(MsgDefParseError::InvalidFormat)?;
                        let array_size_node = type_kind_node.named_child(1);
                        let base_type = Self::parse_base_type(content, base_type_node)?;
                        let array_size = if let Some(array_size_node) = array_size_node {
                            let array_size = array_size_node.utf8_text(content.as_bytes())
                                .map_err(|_| MsgDefParseError::InvalidFormat)?;
                            Some(array_size.parse::<u32>()?)
                        } else {
                            None
                        };
                        let field_type = FieldType {
                            base_type,
                            array_size,
                            is_array: true,
                        };
                        msg_def.fields.push(Field {
                            field_type,
                            field_name: field_name.to_string(),
                        });
                    },
                    _ => {
                        return Err(MsgDefParseError::InvalidFormat);
                    }
                }
            } else if node.kind() == "constant" {
                // ignore constant definition
            }
            maybe_node = node.next_named_sibling();
        }
    
        Ok(msg_def)
    }

    fn parse_base_type(content: &str, base_type_node: Node) -> Result<BaseType, MsgDefParseError> {
        let pkg_name_node = base_type_node.child_by_field_name("pkgname");
        let _msg_name_node = base_type_node.child_by_field_name("msgname")
            .ok_or(MsgDefParseError::InvalidFormat)?;
        let base_type = base_type_node.utf8_text(content.as_bytes())
            .map_err(|_| MsgDefParseError::InvalidFormat)?;
        let sep = if let Some(pkg_name_node) = pkg_name_node {
            Some(pkg_name_node.end_byte() - pkg_name_node.start_byte())
        } else {
            None
        };
        Ok(BaseType {
            string: base_type.to_string(),
            sep,
        })
    }

    pub fn resolve_field_types(&mut self) {
        let mut msg2pkg = BTreeMap::new();
        for (pkg_name, submap) in self.map.iter_mut() {
            for msg_name in submap.keys() {
                match msg2pkg.entry(msg_name.to_string()) {
                    Entry::Vacant(vacant_entry) => {
                        vacant_entry.insert(vec![pkg_name.to_string()]);
                    },
                    Entry::Occupied(mut occupied_entry) => {
                        occupied_entry.get_mut().push(pkg_name.to_string());
                    },
                }
            }
        }

        for msg_def in self.map.values_mut().map(|m| m.values_mut()).flatten() {
            for base_type in msg_def.fields.iter_mut().map(|f| &mut f.field_type.base_type) {
                if base_type.pkg_name().is_none() && !base_type.is_builtin() {
                    if let Some(pkg_names) = msg2pkg.get(base_type.msg_name()) {
                        if pkg_names.len() == 1 {
                            base_type.string = format!("{}/{}", pkg_names[0], base_type.msg_name());
                            base_type.sep = Some(pkg_names[0].len());
                        }
                    }
                }
            }
        }
    }
}

impl BaseType {
    pub fn from_msg_name(msg_name: String) -> Self {
        Self {
            string: msg_name,
            sep: None,
        }
    }

    pub fn pkg_name(&self) -> Option<&str> {
        self.sep.map(|sep| &self.string[..sep])
    }

    pub fn msg_name(&self) -> &str {
        if let Some(sep) = self.sep {
            &self.string[sep+1..]
        } else {
            &self.string
        }
    }

    pub fn is_builtin(&self) -> bool {
        self.sep.is_none()
        && matches!(
            self.string.as_str(),
            "bool" | "byte" | "char" | "float32" | "float64"
            | "int8" | "uint8" | "int16" | "uint16" | "int32" | "uint32" | "int64" | "uint64"
            | "string" | "time" | "duration"
        )
    }
}

#[cfg(test)]
mod test {
    use crate::msg_def::{BaseType, Field, FieldType, MsgDef, MsgDefMap};

    #[test]
    fn test_parse_ros_msg_def() {
        let message = "\
            string data\n\
            int32[] data_array\n\
            bool is_valid\n\
            std_msgs/Header header\n";
        assert_eq!(
            MsgDefMap::parse_ros_msg_def(message),
            Ok(MsgDef {
                fields: vec![
                    Field {
                        field_type: FieldType {
                            base_type: BaseType::from_msg_name("string".to_string()),
                            array_size: None,
                            is_array: false,
                        },
                        field_name: "data".to_string(),
                    },
                    Field {
                        field_type: FieldType {
                            base_type: BaseType::from_msg_name("int32".to_string()),
                            array_size: None,
                            is_array: true,
                        },
                        field_name: "data_array".to_string(),
                    },
                    Field {
                        field_type: FieldType {
                            base_type: BaseType::from_msg_name("bool".to_string()),
                            array_size: None,
                            is_array: false,
                        },
                        field_name: "is_valid".to_string(),
                    },
                    Field {
                        field_type: FieldType {
                            base_type: BaseType {
                                string: "std_msgs/Header".to_string(),
                                sep: Some(8),
                            },
                            array_size: None,
                            is_array: false,
                        },
                        field_name: "header".to_string(),
                    },
                ],
            })
        );
    }
}
